package main

import (
    "encoding/json"
    "time"

    "github.com/bytedance/Elkeid/plugins/collector/engine"
    "github.com/bytedance/Elkeid/plugins/collector/process"
    plugins "github.com/bytedance/plugins"
    "github.com/mitchellh/mapstructure"
    "go.uber.org/zap"
)

type ProcessHandler struct{}

func (h ProcessHandler) Name() string {
    return "process"
}
func (h ProcessHandler) DataType() int {
    return 5050
}

func (h *ProcessHandler) Handle(c *plugins.Client, cache *engine.Cache, seq string) {
    procs, err := process.Processes(false)
    if err != nil {
        zap.S().Error("Failed to get processes:", err)
    } else {
        zap.S().Infof("Successfully retrieved %d processes", len(procs)) // 打印获取到的进程数量
        for _, p := range procs {
            time.Sleep(process.TraversalInterval)
            cmdline, err := p.Cmdline()
            if err != nil {
                zap.S().Errorf("Failed to get cmdline for process %s: %v", p.Pid(), err)
                continue
            }
            stat, err := p.Stat()
            if err != nil {
                zap.S().Errorf("Failed to get stat for process %s: %v", p.Pid(), err)
                continue
            }
            status, err := p.Status()
            if err != nil {
                zap.S().Errorf("Failed to get status for process %s: %v", p.Pid(), err)
                continue
            }
            ns, _ := p.Namespaces()

            // 获取进程的 Cwd、ExeChecksum、ExeHash 和 Exe
            cwd, err := p.Cwd()
            if err != nil {
                zap.S().Errorf("Failed to get cwd for process %s: %v", p.Pid(), err)
                continue
            }
            checksum, err := p.ExeChecksum()
            if err != nil {
                zap.S().Errorf("Failed to get checksum for process %s: %v", p.Pid(), err)
                continue
            }
            exeHash, err := p.ExeHash()
            if err != nil {
                zap.S().Errorf("Failed to get exe hash for process %s: %v", p.Pid(), err)
                continue
            }
            exe, err := p.Exe()
            if err != nil {
                zap.S().Errorf("Failed to get exe for process %s: %v", p.Pid(), err)
                continue
            }

            // 将进程信息转换为 JSON
            processInfo := map[string]interface{}{
                "cmdline":  cmdline,
                "cwd":      cwd,
                "checksum": checksum,
                "exe_hash": exeHash,
                "exe":      exe,
                "pid":      p.Pid(),
                "stat":     stat,
                "status":   status,
                "ns":       ns,
                "container_type": p.ContainerType,
            }
            processJSON, err := json.Marshal(processInfo)
            if err != nil {
                zap.S().Errorf("Failed to marshal process info to JSON for process %s: %v", p.Pid(), err)
            } else {
                // 打印 JSON 格式的进程信息
                zap.S().Infof("Process info:", string(processJSON))
            }

            rec := &plugins.Record{
                DataType:  int32(h.DataType()),
                Timestamp: time.Now().Unix(),
                Data: &plugins.Payload{
                    Fields: make(map[string]string, 40),
                },
            }
            rec.Data.Fields["cmdline"] = cmdline
            rec.Data.Fields["cwd"] = cwd
            rec.Data.Fields["checksum"] = checksum
            rec.Data.Fields["exe_hash"] = exeHash
            rec.Data.Fields["exe"] = exe
            rec.Data.Fields["pid"] = p.Pid()
            rec.Data.Fields["container_type"] = p.ContainerType
            mapstructure.Decode(stat, &rec.Data.Fields)
            mapstructure.Decode(status, &rec.Data.Fields)
            mapstructure.Decode(ns, &rec.Data.Fields)
            m, _ := cache.Get(5056, ns.Pid)
            rec.Data.Fields["container_id"] = m["container_id"]
            rec.Data.Fields["container_name"] = m["container_name"]
            rec.Data.Fields["integrity"] = "true"
            // only for host files
            if _, ok := cache.Get(5057, rec.Data.Fields["exe"]); ok && rec.Data.Fields["container_id"] == "" {
                rec.Data.Fields["integrity"] = "false"
            }
            rec.Data.Fields["package_seq"] = seq
            c.SendRecord(rec)
        }
    }
}

